#include <opencv2/features2d.hpp>
#include <opencv2/videoio.hpp>
#include <opencv2/opencv.hpp>
#include <opencv2/highgui.hpp>      //for imshow
#include <vector>
#include <iostream>
#include <iomanip>

#include "stats.h" // Stats structure definition
#include "utils.h" // Drawing and printing functions

using namespace std;
using namespace cv;

const double akaze_thresh = 3e-4; // AKAZE detection threshold set to locate about 1000 keypoints
const double ransac_thresh = 2.5f; // RANSAC inlier threshold
const double nn_match_ratio = 0.8f; // Nearest-neighbour matching ratio
const int bb_min_inliers = 10; // Minimal number of inliers to draw bounding box
const int stats_update_period = 10; // On-screen statistics are updated every 10 frames

namespace example {
class Tracker
{
public:
    Tracker(const vector<KeyPoint>& kp, const Mat& desc, vector<Point2f> bound, Ptr<DescriptorMatcher> _matcher) :
        matcher(_matcher),
        first_kp(kp),
        first_desc(desc),
        object_bb(bound)
    {
    }

    Mat process(const Mat frame, Stats& stats, const vector<KeyPoint>& kp, const Mat& desc);

protected:
    Ptr<DescriptorMatcher> matcher;

    Mat first_frame, first_desc;
    vector<KeyPoint> first_kp;
    vector<Point2f> object_bb;
};

Mat Tracker::process(const Mat frame, Stats& stats, const vector<KeyPoint>& kp, const Mat& desc)
{
    vector< vector<DMatch> > matches;
    vector<KeyPoint> matched1, matched2;
    matcher->knnMatch(first_desc, desc, matches, 2);
    for(unsigned i = 0; i < matches.size(); i++) {
        if(matches[i][0].distance < nn_match_ratio * matches[i][1].distance) {
            matched1.push_back(first_kp[matches[i][0].queryIdx]);
            matched2.push_back(      kp[matches[i][0].trainIdx]);
        }
    }
    stats.matches = (int)matched1.size();
    Mat inlier_mask, homography;
    vector<KeyPoint> inliers1, inliers2;
    vector<DMatch> inlier_matches;
    if(matched1.size() >= 4) {
        homography = findHomography(Points(matched1), Points(matched2), RANSAC, ransac_thresh, inlier_mask);
    }

    Mat frame_with_bb = frame.clone();

    if(matched1.size() < 4 || homography.empty()) {
        stats.inliers = 0;
        stats.ratio = 0;
        return frame_with_bb;
    }

    for(unsigned i = 0; i < matched1.size(); i++) {
        if(inlier_mask.at<uchar>(i)) {
            int new_i = static_cast<int>(inliers1.size());
            inliers1.push_back(matched1[i]);
            inliers2.push_back(matched2[i]);
            inlier_matches.push_back(DMatch(new_i, new_i, 0));
        }
    }
    stats.inliers = (int)inliers1.size();
    stats.ratio = stats.inliers * 1.0 / stats.matches;

    vector<Point2f> new_bb;
    perspectiveTransform(object_bb, new_bb, homography);

    printf("inliers %d\n", stats.inliers);

    if(stats.inliers >= bb_min_inliers) {
        printf("inliers ------------\n");
        drawBoundingBox(frame_with_bb, new_bb);
    }
    return frame_with_bb;
}
}

int main(int argc, char **argv)
{
    if (argc < 2) {
        cerr << "Usage: " << endl
             << "akaze_track input_path" << endl
             << "  (input_path can be a camera id, like 0,1,2 or a video filename)" << endl;
        return 1;
    }

    std::string video_name = argv[1];
    std::stringstream ssFormat;
    ssFormat << atoi(argv[1]);

    VideoCapture video_in;
    if (video_name.compare(ssFormat.str())==0) {    //test str==str(num)
        video_in.open(atoi(argv[1]));
    }
    else {
        video_in.open(video_name);
    }

    if(!video_in.isOpened()) {
        cerr << "Couldn't open " << argv[1] << endl;
        return 1;
    }

    Stats stats, akaze_stats;
    vector<KeyPoint> kp;
    Mat desc;

    // 构造akaze
    Ptr<AKAZE> akaze = AKAZE::create();
    akaze->setThreshold(akaze_thresh);

    Mat imgFrame = imread("object.jpg");
    akaze->detectAndCompute(imgFrame, noArray(), kp, desc);

    vector<Point2f> bound(4);
    bound[0] = Point(0, 0);
    bound[1] = Point(imgFrame.cols, 0);  
    bound[2] = Point(imgFrame.cols, imgFrame.rows); 
    bound[3] = Point(0, imgFrame.rows);  

    // 暴力匹配
    Ptr<DescriptorMatcher> matcher = DescriptorMatcher::create("BruteForce-Hamming");

    // 构造tracker
    example::Tracker akaze_tracker(kp, desc, bound, matcher);

    namedWindow(video_name, WINDOW_NORMAL);

    Stats akaze_draw_stats;
    int i = 0;

    Mat frame;
    for(;;) {
        i++;
        bool update_stats = (i % stats_update_period == 0);
        video_in >> frame;
        // stop the program if no more images
        if(frame.empty()) break;

        // 客户端
        vector<KeyPoint> kp;
        Mat desc;
        akaze->detectAndCompute(frame, noArray(), kp, desc);
        stats.keypoints = (int)kp.size();

        // 服务端
        Mat akaze_res = akaze_tracker.process(frame, stats, kp, desc);
        akaze_stats += stats;
        if(update_stats) {
            akaze_draw_stats = stats;
        }

        // 展示信息
        drawStatistics(akaze_res, akaze_draw_stats);
        cv::imshow(video_name, akaze_res);

        if (waitKey(1) == 27) break; //quit on ESC button
    }
    akaze_stats /= i - 1;
    printStatistics("AKAZE", akaze_stats);
    return 0;
}
